import pickle
import sys

import networkx as nx
import numpy as np
import scipy
import scipy.sparse as sp


# 定义一个函数load_data，默认参数prefix设为'DBLP'，用于指定数据集的前缀（如DBLP、IMDB等）。
def load_data(prefix='DBLP'):
    from data_loader import data_loader
    # 使用data_loader加载以prefix为路径的一组数据集，'data/'+prefix表示数据集文件夹的路径，例如'data/DBLP'，并将结果存储在dl变量中。dl是一个包含数据和元数据信息的对象。
    dl = data_loader('data/' + prefix)
    # 初始化一个空列表features，用于存储每种节点类型的特征矩阵。
    features = []
    # 遍历节点类型的数量（dl.nodes['count']存储了每种节点类型的节点数量）。例如，如果图中有两种节点类型，则dl.nodes['count']的长度为2，循环两次。
    for i in range(len(dl.nodes['count'])):
        # 从dl.nodes['attr']中获取第i种节点类型的特征矩阵th，这可能是一个稀疏矩阵或其他格式。如果没有特征，则th为None。
        th = dl.nodes['attr'][i]
        # 如果th是None（表示该节点类型没有特征），则为该节点类型创建一个单位矩阵（用scipy的eye()函数生成，维度为该节点类型的节点数量），表示该节点类型的特征就是节点的身份矩阵。
        if th is None:
            features.append(sp.eye(dl.nodes['count'][i]))
        # 果th不为None，则直接将该特征矩阵添加到features列表中。
        else:
            features.append(th)
    # 将所有类型的链接（即图中的边）加总得到整个图的邻接矩阵adjM。dl.links['data']包含了每种关系类型的邻接矩阵，将它们相加得到最终的图邻接矩阵。
    # sum() 函数会将这些邻接矩阵进行逐项相加。
    # 各个位置的元素是这些邻接矩阵相应位置元素相加的结果。例如：
    # relation_1[0, 1] = 1
    # relation_2[0, 1] = 0
    # relation_3[0, 1] = 0
    # 相加后得到 adjM[0, 1] = 1。
    # 最终得到一个不分链接类型的邻接矩阵
    adjM = sum(dl.links['data'].values())
    # 初始化一个全零矩阵labels，大小为第0类节点的数量（dl.nodes['count'][0]）乘以训练集中标签类别的数量（dl.labels_train['num_classes']）。
    # labels用于存储每个节点的标签（one-hot编码）。维度为 (节点类型0的数量, 类别数量)（二维数组）
    # 某些任务（如分类或标签分配）只针对某一类型的节点执行。在这种情况下，labels 只需要为特定类型（这里是类型 0）的节点创建。类型 0 的节点可能是图中的关键节点类型，如用户、文章等，而其他类型的节点可能与这个任务无关。
    # eg：DBLP中类型0是Author，我们的任务就是对Author作分类
    labels = np.zeros((dl.nodes['count'][0], dl.labels_train['num_classes']), dtype=int)
    # 定义验证集的比例为20%，即训练集中20%的数据将用作验证集。
    val_ratio = 0.2
    # 获取训练集中有效标签（非零）的索引，dl.labels_train['mask']是一个布尔掩码，表示哪些节点有标签，np.nonzero()返回这些节点的索引。
    # [0]：np.nonzero 返回的结果是一个元组，其中第一个元素就是我们需要的一维数组中非零元素的索引。[0] 提取了这个一维索引数组。
    train_idx = np.nonzero(dl.labels_train['mask'])[0]
    # 打乱训练集中节点的索引顺序，确保训练集的划分是随机的。
    np.random.shuffle(train_idx)
    # 计算验证集的大小，将训练集中的前split个节点划分为验证集，train_idx.shape[0]表示训练集中节点的数量。
    split = int(train_idx.shape[0] * val_ratio)
    # 将训练集的前split个节点作为验证集的索引val_idx，剩下的节点作为新的训练集的索引train_idx。
    val_idx = train_idx[:split]
    train_idx = train_idx[split:]
    # 对训练集和验证集的索引进行排序，确保索引是按升序排列的。
    train_idx = np.sort(train_idx)
    val_idx = np.sort(val_idx)
    # 获取测试集的节点索引，dl.labels_test['mask']是一个布尔掩码，表示哪些节点属于测试集。
    test_idx = np.nonzero(dl.labels_test['mask'])[0]
    # 为训练集和验证集的节点分配标签。dl.labels_train['data']包含了训练集中节点的标签，按train_idx和val_idx分别将这些标签分配到labels矩阵中。
    labels[train_idx] = dl.labels_train['data'][train_idx]
    labels[val_idx] = dl.labels_train['data'][val_idx]
    # 如果数据集不是IMDB或IMDB-HGB，则对标签进行处理，将one-hot编码的标签转换为类别编号（即对每个节点选择最大的那个one-hot值所对应的类别）。
    if prefix != 'IMDB' and prefix != 'IMDB-HGB':
        labels = labels.argmax(axis=1)
    # 一个字典train_val_test_idx，用于存储训练集、验证集和测试集的节点索引。
    train_val_test_idx = {'train_idx': train_idx, 'val_idx': val_idx, 'test_idx': test_idx}
    # 将训练集、验证集和测试集的节点索引分别存储在字典train_val_test_idx中。
    # 返回5个值：
    # features: 包含所有节点类型特征的列表。
    # adjM: 图的邻接矩阵。
    # labels: 所有节点的标签（训练集、验证集、测试集）。
    # train_val_test_idx: 一个字典，包含训练集、验证集和测试集的索引。
    # dl: 数据加载器对象，保存原始的图数据和标签信息。
    return features, \
        adjM, \
        labels, \
        train_val_test_idx, \
        dl
